import math
import torch
import torch.nn as nn

MASK_VALUE = -1e9

class CrossAttention(nn.Module):
    def __init__(self, sequence_dim=1280, structure_dim=1024, num_heads=8):
        super().__init__()
        self.num_heads = num_heads
        self.sequence_dim = sequence_dim
        self.structure_dim = structure_dim

        self.q_linear = nn.Linear(structure_dim, sequence_dim)
        self.k_linear = nn.Linear(sequence_dim, sequence_dim)
        self.v_linear = nn.Linear(sequence_dim, sequence_dim)

        self.out = nn.Linear(sequence_dim, sequence_dim)

    def forward(self, sequence_feature, structure_feature, sequence_mask=None, structure_mask=None):
        """
        Args:
            sequence_feature (torch.Tensor): Tensor of shape [B, L, sequence_dim] containing the sequence features.
            structure_feature (torch.Tensor): Tensor of shape [B, L, structure_dim] containing the structure features.
            sequence_mask (torch.Tensor, optional): Tensor of shape [B, L] serving as a mask for the sequence features.
            structure_mask (torch.Tensor, optional): Tensor of shape [B, L] serving as a mask for the structure features.

        Returns:
            output (torch.Tensor): Tensor of shape [B, L, sequence_dim] containing the output features.
        """
        B, L, _ = sequence_feature.size()

        q = self.q_linear(structure_feature).view(B, L, self.num_heads, self.sequence_dim // self.num_heads).transpose(1, 2)
        if structure_mask is not None:
            q = q.masked_fill(structure_mask.unsqueeze(1).unsqueeze(-1) == 0, MASK_VALUE)

        k = self.k_linear(sequence_feature).view(B, L, self.num_heads, self.sequence_dim // self.num_heads).transpose(1, 2)
        v = self.v_linear(sequence_feature).view(B, L, self.num_heads, self.sequence_dim // self.num_heads).transpose(1, 2)

        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.sequence_dim ** 0.5)
        if sequence_mask is not None:
            scores = scores.masked_fill(sequence_mask.unsqueeze(1).unsqueeze(-1) == 0, MASK_VALUE)

        attention = torch.softmax(scores, dim=-1)
        context = torch.matmul(attention, v)
        context = context.transpose(1, 2).contiguous().view(B, L, -1)

        output = self.out(context)

        return output
    

class TransformerLayer(nn.Module):
    def __init__(self, sequence_dim=1280, structure_dim=1024, num_heads=8, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(sequence_dim)
        self.norm2 = nn.LayerNorm(sequence_dim)
        self.cross_attention = CrossAttention(sequence_dim, structure_dim, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(sequence_dim, 4 * sequence_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(4 * sequence_dim, sequence_dim),
        )
        self.dropout = nn.Dropout(dropout)

    def forward(self, sequence_feature, structure_feature, sequence_mask=None, structure_mask=None):
        """
        Args:
            sequence_feature (torch.Tensor): Tensor of shape [B, L, sequence_dim] containing the sequence features.
            structure_feature (torch.Tensor): Tensor of shape [B, L, structure_dim] containing the structure features.
            sequence_mask (torch.Tensor, optional): Tensor of shape [B, L] serving as a mask for the sequence features.
            structure_mask (torch.Tensor, optional): Tensor of shape [B, L] serving as a mask for the structure features.

        Returns:
            output (torch.Tensor): Tensor of shape [B, L, sequence_dim] containing the output features.
        """
        sequence_feature = self.norm1(sequence_feature + self.dropout(self.cross_attention(sequence_feature, structure_feature, sequence_mask, structure_mask)))
        sequence_feature = self.norm2(sequence_feature + self.dropout(self.feed_forward(sequence_feature)))

        return sequence_feature
    

class PositionalEncoding(nn.Module):
    def __init__(self, sequence_dim, max_len=2000):
        super(PositionalEncoding, self).__init__()
        
        pe = torch.zeros(max_len, sequence_dim)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, sequence_dim, 2) *
                             -(math.log(10000.0) / sequence_dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return x


class Transformer(nn.Module):
    def __init__(self, sequence_dim=1280, structure_dim=1024, num_heads=8, num_layers=6):
        super().__init__()
        self.sequence_pos_encoder = PositionalEncoding(sequence_dim)
        self.structure_pos_encoder = PositionalEncoding(structure_dim)
        self.encoder = TransformerEncoder(sequence_dim, structure_dim, num_heads, num_layers)
        self.decoder = TransformerDecoder(sequence_dim, structure_dim, num_heads, num_layers)

    def forward(self, sequence_feature, structure_feature, sequence_mask=None, structure_mask=None):
        """
        Args:
            sequence_feature (torch.Tensor): Tensor of shape [B, L, sequence_dim] containing the sequence features.
            structure_feature (torch.Tensor): Tensor of shape [B, L, structure_dim] containing the structure features.
            sequence_mask (torch.Tensor, optional): Tensor of shape [B, L] serving as a mask for the sequence features.
            structure_mask (torch.Tensor, optional): Tensor of shape [B, L] serving as a mask for the structure features.

        Returns:
            output (torch.Tensor): Tensor of shape [B, L, sequence_dim] containing the output features.
        """
        sequence_feature = self.sequence_pos_encoder(sequence_feature)
        structure_feature = self.structure_pos_encoder(structure_feature)
        encoder_output = self.encoder(sequence_feature, structure_feature, sequence_mask, structure_mask)
        output = self.decoder(encoder_output, structure_feature, sequence_mask, structure_mask)

        return output


class TransformerEncoder(nn.Module):
    def __init__(self, sequence_dim=1280, structure_dim=1024, num_heads=8, num_layers=6):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerLayer(sequence_dim, structure_dim, num_heads)
            for _ in range(num_layers)
        ])

    def forward(self, sequence_feature, structure_feature, sequence_mask=None, structure_mask=None):
        for layer in self.layers:
            sequence_feature = layer(sequence_feature, structure_feature, sequence_mask, structure_mask)

        return sequence_feature


class TransformerDecoder(nn.Module):
    def __init__(self, sequence_dim=1280, structure_dim=1024, num_heads=8, num_layers=6):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerLayer(sequence_dim, structure_dim, num_heads)
            for _ in range(num_layers)
        ])

    def forward(self, sequence_feature, structure_feature, sequence_mask=None, structure_mask=None):
        for layer in self.layers:
            sequence_feature = layer(sequence_feature, structure_feature, sequence_mask, structure_mask)

        return sequence_feature
    

if __name__ == "__main__":
    # Test the CrossAttention module
    B, L, sequence_dim, structure_dim = 2, 10, 1280, 1024
    sequence_feature = torch.randn(B, L, sequence_dim)
    structure_feature = torch.randn(B, L, structure_dim)

    cross_attention = CrossAttention(sequence_dim, structure_dim)
    output = cross_attention(sequence_feature, structure_feature)

    print(f"Output shape: {output.shape}")

    # Test the Transformer model
    model = Transformer(sequence_dim=sequence_dim, structure_dim=structure_dim, num_heads=8, num_layers=6)
    output = model(sequence_feature, structure_feature)
    print(f"Transformer Output shape: {output.shape}")